import os
import numpy as np
import pandas as pd
import csv
import glob
import sys
import importlib
import traceback
import matplotlib.pyplot as plt
from matplotlib.font_manager import FontProperties  # 新增：字体管理
from scipy.optimize import curve_fit
from scipy.linalg import eigh
from concurrent.futures import ProcessPoolExecutor, as_completed
import multiprocessing as mp

# 添加当前目录到系统路径
sys.path.append(os.path.dirname(os.path.abspath(__file__)))

# 配置中文字体（解决中文显示问题）
plt.rcParams["font.family"] = ["SimHei", "WenQuanYi Micro Hei", "Heiti TC"]
plt.rcParams['axes.unicode_minus'] = False  # 解决负号显示问题
font_prop = FontProperties(family=["SimHei", "WenQuanYi Micro Hei", "Heiti TC"])

# 动态导入必要的模块和函数
try:
    module_3d_pixel = importlib.import_module('3D_pixel')
    SimulationConfig = getattr(module_3d_pixel, 'SimulationConfig')
    read_deposits_info = getattr(module_3d_pixel, 'read_deposits_info')
    process_deposits = getattr(module_3d_pixel, 'process_deposits')
    sample_trajectory = getattr(module_3d_pixel, 'sample_trajectory')
    voxelize_electrons = getattr(module_3d_pixel, 'voxelize_electrons')
except ImportError as e:
    print(f"导入模块失败: {e}")
    sys.exit(1)


# ---------------------- 三维自适应图像矩计算器（优化版） ----------------------
class ImageMomentCalculator:
    def __init__(self, config):
        self.config = config
        self.voxel_size = np.array(config.voxel_size)
        # 优化1：提高最小体素数阈值（过滤更短径迹）
        self.min_voxels = 8  # 从5提升至8，确保径迹有足够几何信息
        # 优化2：添加径迹长度阈值（单位：mm）
        self.min_track_length = 0.5  # 过滤短于0.5mm的径迹
        # 优化3：三阶矩筛选阈值（用于判断头部对称性）
        self.max_third_moment_norm = 0.3  # 三阶矩模长上限

    def _get_physical_coords(self, voxels, offset):
        return (voxels + offset) * self.voxel_size

    def calculate_3d_moments(self, voxel_grid, offset):
        voxels = np.argwhere(voxel_grid > 0)
        if len(voxels) < self.min_voxels:
            return None, None, None, voxels
        
        coords = self._get_physical_coords(voxels, offset)
        densities = voxel_grid[voxel_grid > 0]
        total_mass = np.sum(densities)
        
        # 优化：加权重心计算时增加密度权重的平方（增强强沉积点的影响）
        weights = densities ** 1.2  # 权重指数从1.0调整为1.2
        weights /= np.sum(weights)  # 归一化权重
        com = np.sum(coords * weights[:, np.newaxis], axis=0)
        
        centered_coords = coords - com
        # 优化：协方差矩阵计算使用加权平均（更稳健）
        cov_matrix = np.cov(centered_coords.T, aweights=densities)
        
        # 三阶矩计算（带密度权重）
        densities_expanded = np.expand_dims(densities, axis=1)
        third_moment = np.sum(densities_expanded * np.power(centered_coords, 3), axis=0) / total_mass
        
        return com, cov_matrix, third_moment, voxels

    def get_principal_direction(self, cov_matrix):
        eigenvalues, eigenvectors = eigh(cov_matrix, eigvals=(2, 2))
        main_dir = eigenvectors.flatten()
        return main_dir / np.linalg.norm(main_dir)

    def split_by_plane(self, coords, densities, main_dir, com, distance):
        projections = np.dot(coords - com, main_dir)
        mask = projections <= distance
        return coords[mask], densities[mask], coords[~mask], densities[~mask]

    def adaptive_cut(self, coords, densities, main_dir, com):
        # 优化4：增加滑动平面数量（从20增至50，提高切割精度）
        min_required_voxels = max(self.min_voxels, int(len(coords) * 0.3))  # 至少保留30%的体素
        if len(coords) < min_required_voxels * 2:
            return coords, densities
        
        projections = np.dot(coords - com, main_dir)
        min_proj, max_proj = np.min(projections), np.max(projections)
        # 滑动平面在有效范围内加密（跳过两端10%的极端区域）
        slide_range = np.linspace(min_proj + 0.1*(max_proj-min_proj), 
                                 max_proj - 0.1*(max_proj-min_proj), 50)
        
        best_cut_pos = max_proj
        best_moment_norm = np.inf
        
        for pos in slide_range:
            head_coords, head_dens, _, _ = self.split_by_plane(
                coords, densities, main_dir, com, pos
            )
            if len(head_coords) < min_required_voxels:
                continue
            
            # 优化：头部重心计算使用更高权重
            head_weights = head_dens ** 1.2
            head_weights /= np.sum(head_weights)
            head_com = np.sum(head_coords * head_weights[:, np.newaxis], axis=0)
            
            head_centered = head_coords - head_com
            head_dens_expanded = np.expand_dims(head_dens, axis=1)
            head_third_moment = np.sum(head_dens_expanded * np.power(head_centered, 3), axis=0) / np.sum(head_dens)
            moment_norm = np.linalg.norm(head_third_moment)
            
            if moment_norm < best_moment_norm:
                best_moment_norm = moment_norm
                best_cut_pos = pos
        
        head_coords, head_dens, _, _ = self.split_by_plane(
            coords, densities, main_dir, com, best_cut_pos
        )
        return head_coords, head_dens

    def calculate_polarization_angles(self, main_dir):
        theta = np.arccos(np.clip(main_dir[2], -1.0, 1.0))
        phi = np.arctan2(main_dir[1], main_dir[0])
        phi = phi + 2 * np.pi if phi < 0 else phi  # 转换为0-2π
        return theta, phi

    def analyze_track(self, voxel_grid, offset, visualize=False):
        # 第一步：计算初始矩并进行初步筛选
        com_initial, cov_initial, third_moment_initial, voxels = self.calculate_3d_moments(voxel_grid, offset)
        if com_initial is None:
            return None
        
        # 优化5：计算初始径迹长度并筛选
        coords = self._get_physical_coords(voxels, offset)
        track_length = np.linalg.norm(np.max(coords, axis=0) - np.min(coords, axis=0))
        if track_length < self.min_track_length:
            return None  # 过滤过短径迹
        
        # 优化6：初始三阶矩筛选（过于不对称的径迹直接剔除）
        if np.linalg.norm(third_moment_initial) > self.max_third_moment_norm * 1.5:
            return None
        
        main_dir_initial = self.get_principal_direction(cov_initial)
        densities = voxel_grid[voxel_grid > 0]
        
        # 自适应切割
        head_coords, head_dens = self.adaptive_cut(coords, densities, main_dir_initial, com_initial)
        if len(head_coords) < self.min_voxels:
            return None
        
        # 优化7：头部区域质量筛选
        head_length = np.linalg.norm(np.max(head_coords, axis=0) - np.min(head_coords, axis=0))
        if head_length < self.min_track_length * 0.5:  # 头部长度至少为总长度的50%
            return None
        
        # 计算头部矩特征
        head_weights = head_dens ** 1.2
        head_weights /= np.sum(head_weights)
        head_com = np.sum(head_coords * head_weights[:, np.newaxis], axis=0)
        head_centered = head_coords - head_com
        head_cov = np.cov(head_centered.T, aweights=head_dens)
        head_main_dir = self.get_principal_direction(head_cov)
        
        # 优化8：头部三阶矩最终筛选
        head_total_mass = np.sum(head_dens)
        head_third_moment = np.sum(np.expand_dims(head_dens, axis=1) * np.power(head_centered, 3), axis=0) / head_total_mass
        if np.linalg.norm(head_third_moment) > self.max_third_moment_norm:
            return None
        
        polarization_theta, polarization_phi = self.calculate_polarization_angles(head_main_dir)
        eigenvalues = np.sort(eigh(head_cov, eigvals_only=True))[::-1]
        
        # 优化9：通过特征值比筛选（排除过于扁平的径迹）
        if eigenvalues[1] / eigenvalues[0] > 0.8:  # 第二特征值与第一特征值比过高，方向不明显
            return None
        
        initial_results = {
            'center_of_mass': com_initial,
            'main_direction': main_dir_initial,
            'z_start': np.min(coords[:, 2]),
            'z_end': np.max(coords[:, 2]),
            'length': track_length
        }
        final_results = {
            'center_of_mass': head_com,
            'main_direction': head_main_dir,
            'z_start': np.min(head_coords[:, 2]),
            'z_end': np.max(head_coords[:, 2]),
            'eigenvalues': eigenvalues,
            'length': head_length
        }
        
        return {
            'initial_results': initial_results,
            'final_results': final_results,
            'polarization_theta': polarization_theta,
            'polarization_phi': polarization_phi,
        }


# ---------------------- 全局并行处理函数 ----------------------
def process_single_file_worker(csv_file, output_table_dir, output_image_dir):
    try:
        file_name = os.path.basename(csv_file)
        event_name = os.path.splitext(file_name)[0]
        # 减少打印频率，避免输出刷屏
        # print(f"处理文件：{file_name}")
        
        deposits = read_deposits_info(csv_file)
        if not deposits:
            return None, event_name
            
        config = SimulationConfig()
        charge_positions = process_deposits(deposits, config)
        if len(charge_positions) == 0:
            return None, event_name
            
        sampled_positions = sample_trajectory(charge_positions, config)
        if len(sampled_positions) == 0:
            return None, event_name
        
        voxel_grid, offset = voxelize_electrons(sampled_positions[:, :3], config)
        active_voxels = np.sum(voxel_grid > 0)
        if active_voxels < 8:  # 更严格的体素数筛选
            return None, event_name
            
        moment_calculator = ImageMomentCalculator(config)
        results = moment_calculator.analyze_track(voxel_grid, offset, visualize=False)
        
        if results is None:
            return None, event_name
            
        initial_results = results['initial_results']
        final_results = results['final_results']
        
        result_data = {
            'event_name': event_name,
            'initial_com_x': initial_results['center_of_mass'][0],
            'initial_com_y': initial_results['center_of_mass'][1],
            'initial_com_z': initial_results['center_of_mass'][2],
            'initial_main_dir_x': initial_results['main_direction'][0],
            'initial_main_dir_y': initial_results['main_direction'][1],
            'initial_main_dir_z': initial_results['main_direction'][2],
            'initial_track_length': initial_results['length'],
            'final_com_x': final_results['center_of_mass'][0],
            'final_com_y': final_results['center_of_mass'][1],
            'final_com_z': final_results['center_of_mass'][2],
            'final_main_dir_x': final_results['main_direction'][0],
            'final_main_dir_y': final_results['main_direction'][1],
            'final_main_dir_z': final_results['main_direction'][2],
            'final_track_length': final_results['length'],
            'polarization_theta_deg': np.degrees(results['polarization_theta']),
            'polarization_phi_deg': np.degrees(results['polarization_phi']),
            'eigenvalue_1': final_results['eigenvalues'][0],
            'eigenvalue_2': final_results['eigenvalues'][1],
            'eigenvalue_3': final_results['eigenvalues'][2],
        }
        
        output_csv_path = os.path.join(output_table_dir, f"{event_name}_3d_image_dimension.csv")
        df = pd.DataFrame([result_data])
        df.to_csv(output_csv_path, index=False)
        
        phi_deg = np.degrees(results['polarization_phi'])
        return phi_deg, event_name
        
    except Exception as e:
        # 仅打印关键错误
        if "critical" in str(e).lower():
            print(f"处理文件 {file_name} 时出错：{e}")
        return None, os.path.splitext(os.path.basename(csv_file))[0]


# ---------------------- 批量处理器类（优化版） ----------------------
class Batch3DImageDimensionProcessor:
    def __init__(self, input_dir, output_root_dir, max_workers=None):
        self.input_dir = input_dir
        self.output_root_dir = output_root_dir
        self.output_table_dir = os.path.join(output_root_dir, 'tables')
        self.output_image_dir = os.path.join(output_root_dir, 'images')
        self._create_output_directories()
        self.max_workers = max_workers or min(8, mp.cpu_count())
        print(f"使用 {self.max_workers} 个并行工作进程")
    
    def _create_output_directories(self):
        try:
            os.makedirs(self.output_table_dir, exist_ok=True)
            os.makedirs(self.output_image_dir, exist_ok=True)
            print(f"已创建输出目录：")
            print(f"- 表格数据目录：{self.output_table_dir}")
            print(f"- 图像数据目录：{self.output_image_dir}")
        except Exception as e:
            print(f"创建输出目录时出错：{e}")
    
    def find_csv_files(self):
        csv_files = []
        try:
            csv_pattern = os.path.join(self.input_dir, '*.csv')
            csv_files = glob.glob(csv_pattern)
            print(f"在 {self.input_dir} 目录下找到 {len(csv_files)} 个CSV文件")
        except Exception as e:
            print(f"查找CSV文件时出错：{e}")
        return csv_files
    
    def process_all_files_parallel(self):
        csv_files = self.find_csv_files()
        total_files = len(csv_files)
        if total_files == 0:
            print(f"警告：在 {self.input_dir} 目录下未找到任何CSV文件")
            return
        
        all_phi_angles = []
        success_count = 0
        print(f"\n开始并行处理 {total_files} 个文件，使用 {self.max_workers} 个工作进程...")
        
        with ProcessPoolExecutor(max_workers=self.max_workers) as executor:
            future_to_file = {
                executor.submit(process_single_file_worker, csv_file, self.output_table_dir, self.output_image_dir): csv_file 
                for csv_file in csv_files
            }
            
            for i, future in enumerate(as_completed(future_to_file)):
                csv_file = future_to_file[future]
                try:
                    phi_angle, event_name = future.result()
                    if phi_angle is not None:
                        all_phi_angles.append(phi_angle)
                        success_count += 1
                        # 每1000个事件打印一次进度
                        if success_count % 1000 == 0:
                            print(f"已完成 {success_count}/{total_files} 个事件（{success_count/total_files*100:.1f}%）")
                except Exception as e:
                    if i % 1000 == 0:  # 偶尔打印错误
                        print(f"处理文件 {csv_file} 时发生异常: {e}")
        
        print(f"\n==================== 批量处理完成 ====================")
        print(f"总文件数：{total_files}")
        print(f"成功处理：{success_count}")
        print(f"处理失败：{total_files - success_count}")
        print(f"有效事件保留率：{success_count/total_files*100:.1f}%")
        print(f"结果保存在：{self.output_root_dir}")
        
        if all_phi_angles:
            print(f"\n生成φ角分布直方图及调制度拟合...")
            self.generate_phi_distribution_histogram(all_phi_angles)
        else:
            print("没有收集到有效的phi角数据，无法生成直方图")
    
    def generate_phi_distribution_histogram(self, phi_angles):
        if not phi_angles:
            print("没有收集到phi角数据，无法生成直方图")
            return
        
        try:
            # 优化10：过滤phi角异常值（偏离均值±3σ的数据）
            phi_mean = np.mean(phi_angles)
            phi_std = np.std(phi_angles)
            valid_phi = [phi for phi in phi_angles 
                         if (phi_mean - 3*phi_std) <= phi <= (phi_mean + 3*phi_std)]
            print(f"过滤前事件数：{len(phi_angles)}，过滤后事件数：{len(valid_phi)}")
            if len(valid_phi) < len(phi_angles)*0.5:  # 异常值过多时不过滤
                valid_phi = phi_angles
            
            plt.figure(figsize=(10, 7))
            bins = 36
            hist, bin_edges = np.histogram(valid_phi, bins=bins, range=(0, 360))
            bin_centers = (bin_edges[:-1] + bin_edges[1:]) / 2
            
            def cosine_squared(phi, A, B, phi0):
                return A + B * np.square(np.cos((phi - phi0) * np.pi / 180))
            
            # 优化11：更稳健的初始参数估算
            A_init = np.mean(hist)
            B_init = max(0.1, np.max(hist) - A_init)  # 确保B初始值至少为0.1
            phi0_init = np.median(valid_phi) % 360  # 用中位数替代均值，抗干扰更强
            initial_params = [A_init, B_init, phi0_init]
            
            # 优化12：放宽边界容差，避免数值精度问题
            popt, pcov = curve_fit(
                cosine_squared, 
                bin_centers, 
                hist, 
                p0=initial_params,
                bounds=([-1e-3, -1e-3, -1e-3], [np.inf, np.inf, 360+1e-3]),
                maxfev=10000  # 增加迭代次数
            )
            A, B, phi0 = popt
            B = max(0.0, B)  # 确保B非负
            denominator = 2 * A + B
            modulation = B / denominator if denominator > 1e-6 else 0.0
            
            # 计算拟合优度（修正自由度）
            residuals = hist - cosine_squared(bin_centers, A, B, phi0)
            chi2 = np.sum(np.square(residuals)) / max(1, (bins - len(popt)))
            
            # 绘制拟合曲线和直方图
            phi_fit = np.linspace(0, 360, 360)
            plt.plot(phi_fit, cosine_squared(phi_fit, A, B, phi0), 'r-', linewidth=2, 
                     label=f'余弦平方拟合: μ={modulation:.2f}, φ₀={phi0:.1f}°')
            plt.bar(bin_centers, hist, width=360/bins, alpha=0.7, color='blue', edgecolor='black', label='φ角分布')
            
            # 添加拟合参数文本（使用中文字体）
            fit_text = (f'χ²/ndf = {chi2:.2f}\n'
                       f'A = {A:.1f}\n'
                       f'B = {B:.1f}\n'
                       f'φ₀ = {phi0:.1f}°\n'
                       f'调制度 μ = {modulation:.2f}')
            plt.text(0.05, 0.95, fit_text, transform=plt.gca().transAxes,
                     verticalalignment='top', bbox=dict(boxstyle='round', facecolor='wheat', alpha=0.5),
                     fontproperties=font_prop)
            
            plt.xlabel('出射方位角 φ (°)', fontproperties=font_prop, fontsize=12)
            plt.ylabel('事件数', fontproperties=font_prop, fontsize=12)
            plt.title('光电子径迹的重建出射角度分布', fontproperties=font_prop, fontsize=14)
            plt.grid(True, alpha=0.3)
            plt.legend(prop=font_prop)
            
            output_image_path = os.path.join(self.output_image_dir, 'phi_distribution_fit.png')
            plt.tight_layout()
            plt.savefig(output_image_path, dpi=300, bbox_inches='tight')
            plt.close()
            print(f"出射角度分布图已保存到: {output_image_path}")
            
            # 输出统计信息
            print(f"\nφ角分布统计摘要:")
            print(f"  总有效事件数: {len(valid_phi)}")
            print(f"  平均值: {np.mean(valid_phi):.2f}°")
            print(f"  标准差: {np.std(valid_phi):.2f}°")
            print(f"  调制度 μ: {modulation:.2f}")
            print(f"  偏振相位角 φ₀: {phi0:.1f}°")
            print(f"  拟合优度 χ²/ndf: {chi2:.2f}")
            
        except Exception as e:
            print(f"生成直方图时出错：{e}")
            traceback.print_exc()


# ------------------- 程序运行入口 -------------------
if __name__ == "__main__":
    INPUT_DIR = r"/mnt/d/Track_reconstruction/Graduation/Track_data/sim_8kev_50000"
    OUTPUT_ROOT_DIR = r"/home/qht/Graduation/data_1/u"
    
    processor = Batch3DImageDimensionProcessor(INPUT_DIR, OUTPUT_ROOT_DIR)
    processor.process_all_files_parallel()